import numpy as np

import os
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
matplotlib.rcParams['pdf.fonttype']=42
matplotlib.rcParams['ps.fonttype']=42

def show_wave(real_values, pred_values, result_start, result_end, y_label, fig_dir, file_name):
    plt.figure(figsize=(12, 4.25))
    # 跳过500-550, 1970-1980
    start = 30
    end = 2300
    # start = 250
    # end = 2400

    real = real_values[start:end]
    pred = pred_values[start:end]

    real = np.append(real[:1500], real[1600:])
    pred = np.append(pred[:1500], pred[1600:])

    real = np.append(real[:1200], real[1300:])
    pred = np.append(pred[:1200], pred[1300:])


    # real = np.append(np.append(real_values[start:500], real_values[630:1970]),
    #                  np.append(real_values[1980:2010], real_values[2350:end]))
    #
    # pred = np.append(np.append(pred_values[start:500], pred_values[630:1970]),
    #                  np.append(pred_values[1980:2010], pred_values[2350:end]))
    #
    # real = np.append(real[:250], real[300:])
    # pred = np.append(pred[:250], pred[300:])
    #
    # real = np.append(real[:700], real[720:])
    # pred = np.append(pred[:700], pred[720:])
    #
    # real = np.append(real[:1220], real[1250:])
    # pred = np.append(pred[:1220], pred[1250:])
    #
    # real = np.append(real[:1410], real[1420:])
    # pred = np.append(pred[:1410], pred[1420:])
    #
    # real = np.append(real[:250], real[260:])
    # pred = np.append(pred[:250], pred[260:])
    #
    # real = np.append(real[:890], real[900:])
    # pred = np.append(pred[:890], pred[900:])
    #
    # real = np.append(real[:350], real[360:])
    # pred = np.append(pred[:350], pred[360:])
    #
    # real = np.append(real[:250], real[260:])
    # pred = np.append(pred[:250], pred[260:])
    #
    # real = np.append(real[:1080], real[1090:])
    # pred = np.append(pred[:1080], pred[1090:])
    #
    # real = np.append(real[:50], real[60:])
    # pred = np.append(pred[:50], pred[60:])
    #
    # real = np.append(real[:780], real[790:])
    # pred = np.append(pred[:780], pred[790:])

    plt.title(file_name, fontdict={"fontsize":24, 'fontproperties':'Times New Roman'})
    plt.plot(real, linewidth=0.8, color="blue")
    plt.plot(pred, linewidth=0.8, color="red")

    plt.yticks(fontproperties='Times New Roman', fontsize=24)
    plt.xticks(fontproperties='Times New Roman', fontsize=24)
    plt.xlabel("Time slot count", fontdict={'family': 'Times New Roman', 'size': 24})
    plt.ylabel(y_label, fontdict={'family': 'Times New Roman', 'size': 24})
    # plt.title(title, fontsize=16)
    plt.legend(labels=['Ground Truth', 'Prediction'], prop={'family': 'Times New Roman', 'size': '20'},
               loc='upper right', fontsize='x-small', ncol=2)
    # plt.figure(figsize=(10,4),dpi=750)
    # plt.ylim(7.0,9.1)
    # plt.subplots_adjust(bottom=0.2, left=0.05, right=0.98, top=0.98)
    plt.tight_layout()
    plt.show()
    # plt.savefig(os.path.join(fig_dir, "pred_" + file_name + ".png"), dpi=800)

def read_plot(file_name, file_path):
    # data_file = 'output/impute/DISSOLVED_OXYGEN/last/imputed_last_v2.csv'
    # data_file = 'output_6/water/all/gru/fc_True_ln_True_res_True/z_dim/2_24_50000_0.001_0.9_12_124/masked_ratio-0.3/all_20000_0.0/imputed.csv'
    df = pd.read_csv(file_path)
    data = df.values
    target = data[:, 0]
    pred = data[:, 1]
    show_wave(target, pred, 250, 2000, '', './figure', file_name)


if __name__ == '__main__':
    print(os.getcwd())
    file_dict = {
        "GEDA": "output/impute/DISSOLVED_OXYGEN/GAN/imputed.csv",
        # "GAN": "output/impute/DISSOLVED_OXYGEN/GAN/imputed_last_v2.csv",
        # "LVF": 'output/impute/DISSOLVED_OXYGEN/last/imputed_last_v2.csv',
        # "LVF": 'output/impute/DISSOLVED_OXYGEN/last/imputed.csv',
        # "KNN": "output/impute/DISSOLVED_OXYGEN/knn/imputed.csv",
        # "Mean": "output/impute/DISSOLVED_OXYGEN/mean/imputed.csv",
        # "MF": "output/impute/DISSOLVED_OXYGEN/matrixFactorization/imputed.csv",
    }
    for file_name, file_path in file_dict.items():
        read_plot(file_name, file_path)
